from tqdm import tqdm
import network
import utils
import os
import random
import argparse
import numpy as np

from torch.utils import data
from datasets import VOCSegmentation, Cityscapes
from self import SelfDataset
from utils import ext_transforms as et
from metrics import StreamSegMetrics

import torch
import torch.nn as nn
from utils.visualizer import Visualizer

from PIL import Image
import matplotlib
import matplotlib.pyplot as plt
from torchvision.models.segmentation import deeplabv3_resnet101
from eval import Dataset_Seg, evaluate
from torch.autograd import Variable
from setting import TIME_NOW
from torchvision import transforms

import ssl

ssl._create_default_https_context = ssl._create_unverified_context


def get_argparser():
    parser = argparse.ArgumentParser()

    # Datset Options
    parser.add_argument("--data_root", type=str, default='./datasets/data',
                        help="path to Dataset")
    parser.add_argument("--anno_root", type=str, default='./datasets/data/annotations_trainval',
                        help="path to Annotation")
    parser.add_argument("--dataset", type=str, default='cv',
                        choices=['voc', 'cityscapes', 'cv'], help='Name of dataset')
    parser.add_argument("--num_classes", type=int, default=21,
                        help="num classes (default: None)")
    # Deeplab Options
    parser.add_argument("--model", type=str, default='deeplabv3plus_resnet101',
                        choices=['deeplabv3_resnet50', 'deeplabv3plus_resnet50',
                                 'deeplabv3_resnet101', 'deeplabv3plus_resnet101',
                                 'deeplabv3_mobilenet', 'deeplabv3plus_mobilenet'], help='model name')
    parser.add_argument("--separable_conv", action='store_true', default=False,
                        help="apply separable conv to decoder and aspp")
    parser.add_argument("--output_stride", type=int, default=16, choices=[8, 16])

    # Train Options
    parser.add_argument("--test_only", action='store_true', default=False)
    parser.add_argument("--save_val_results", action='store_true', default=False,
                        help="save segmentation results to \"./results\"")
    parser.add_argument("--total_itrs", type=int, default=3000,
                        help="epoch number (default: 3000)")
    parser.add_argument("--lr", type=float, default=0.0001,
                        help="learning rate (default: 0.0001)")
    parser.add_argument("--lr_policy", type=str, default='poly', choices=['poly', 'step'],
                        help="learning rate scheduler policy")
    parser.add_argument("--step_size", type=int, default=10000)
    parser.add_argument("--crop_val", action='store_true', default=False,
                        help='crop validation (default: False)')
    parser.add_argument("--batch_size", type=int, default=4,
                        help='batch size (default: 4)')
    parser.add_argument("--val_batch_size", type=int, default=1,
                        help='batch size for validation (default: 1)')
    parser.add_argument("--crop_size", type=int, default=513)

    parser.add_argument("--ckpt", default='./checkpoints/best_deeplabv3_resnet101_voc_os16.pth', type=str,
                        help="restore from checkpoint")
    parser.add_argument("--continue_training", action='store_true', default=False)

    parser.add_argument("--loss_type", type=str, default='cross_entropy',
                        choices=['cross_entropy', 'focal_loss'], help="loss type (default: False)")
    parser.add_argument("--gpu_id", type=str, default='0',
                        help="GPU ID")
    parser.add_argument("--weight_decay", type=float, default=1e-4,
                        help='weight decay (default: 1e-4)')
    parser.add_argument("--random_seed", type=int, default=1,
                        help="random seed (default: 1)")
    parser.add_argument("--print_interval", type=int, default=10,
                        help="print interval of loss (default: 10)")
    parser.add_argument("--val_interval", type=int, default=1,
                        help="epoch interval for eval (default: 1)")
    parser.add_argument("--download", action='store_true', default=False,
                        help="download datasets")
    # PASCAL VOC Options
    parser.add_argument("--year", type=str, default='2012',
                        choices=['2012_aug', '2012', '2011', '2009', '2008', '2007'], help='year of VOC')
    # Visdom options
    parser.add_argument("--enable_vis", action='store_true', default=False,
                        help="use visdom for visualization")
    parser.add_argument("--vis_port", type=str, default='60004',
                        help='port for visdom')
    parser.add_argument("--vis_env", type=str, default='main',
                        help='env for visdom')
    parser.add_argument("--vis_num_samples", type=int, default=8,
                        help='number of samples for visualization (default: 8)')
    return parser


def get_dataset(opts):
    if opts.dataset == 'cv':
        train_transform = et.ExtCompose([
            et.ExtResize([opts.crop_size, opts.crop_size]),
            # et.ExtRandomCrop(size=(opts.crop_size, opts.crop_size), pad_if_needed=True),
            et.ExtRandomHorizontalFlip(),
            et.ExtToTensor(),
            et.ExtNormalize(mean=[0.485, 0.456, 0.406],
                            std=[0.229, 0.224, 0.225]),
        ])
        if opts.crop_val:
            val_transform = et.ExtCompose([
                et.ExtResize([opts.crop_size, opts.crop_size]),
                # transforms.RandomHorizontalFlip(),  # 随机水平翻转
                # transforms.RandomRotation(degrees=15),  # 随机旋转
                et.ExtToTensor(),
                et.ExtNormalize(mean=[0.485, 0.456, 0.406],
                                std=[0.229, 0.224, 0.225]),
            ])
        else:
            val_transform = et.ExtCompose([
                et.ExtResize([opts.crop_size, opts.crop_size]),
                # transforms.RandomHorizontalFlip(),  # 随机水平翻转
                # transforms.RandomRotation(degrees=15),  # 随机旋转
                et.ExtToTensor(),
                et.ExtNormalize(mean=[0.485, 0.456, 0.406],
                                std=[0.229, 0.224, 0.225]),
            ])
        train_dst = SelfDataset(imgs_dir=opts.data_root, anno_dir=opts.anno_root, split='train',
                                transform=train_transform)
        val_dst = SelfDataset(imgs_dir=opts.data_root, anno_dir=opts.anno_root, split='val', transform=val_transform)
    if opts.dataset == 'voc':
        train_transform = et.ExtCompose([
            # et.ExtResize(size=opts.crop_size),
            et.ExtRandomScale((0.5, 2.0)),
            et.ExtRandomCrop(size=(opts.crop_size, opts.crop_size), pad_if_needed=True),
            et.ExtRandomHorizontalFlip(),
            et.ExtToTensor(),
            # et.ExtNormalize(mean=[0.485, 0.456, 0.406],
            #                 std=[0.229, 0.224, 0.225]),
        ])
        if opts.crop_val:
            val_transform = et.ExtCompose([
                # et.ExtResize(opts.crop_size),
                # et.ExtCenterCrop(opts.crop_size),
                et.ExtToTensor(),
                # et.ExtNormalize(mean=[0.485, 0.456, 0.406],
                #                 std=[0.229, 0.224, 0.225]),
            ])
        else:
            val_transform = et.ExtCompose([
                et.ExtToTensor(),
                # et.ExtNormalize(mean=[0.485, 0.456, 0.406],
                #                 std=[0.229, 0.224, 0.225]),
            ])
        train_dst = VOCSegmentation(root=opts.data_root, year=opts.year,
                                    image_set='train', download=opts.download, transform=train_transform)
        val_dst = VOCSegmentation(root=opts.data_root, year=opts.year,
                                  image_set='val', download=False, transform=val_transform)
    if opts.dataset == 'cityscapes':
        train_transform = et.ExtCompose([
            # et.ExtResize( 512 ),
            et.ExtRandomCrop(size=(opts.crop_size, opts.crop_size)),
            et.ExtColorJitter(brightness=0.5, contrast=0.5, saturation=0.5),
            et.ExtRandomHorizontalFlip(),
            et.ExtToTensor(),
            et.ExtNormalize(mean=[0.485, 0.456, 0.406],
                            std=[0.229, 0.224, 0.225]),
        ])

        val_transform = et.ExtCompose([
            # et.ExtResize( 512 ),
            et.ExtToTensor(),
            et.ExtNormalize(mean=[0.485, 0.456, 0.406],
                            std=[0.229, 0.224, 0.225]),
        ])

        train_dst = Cityscapes(root=opts.data_root,
                               split='train', transform=train_transform)
        val_dst = Cityscapes(root=opts.data_root,
                             split='val', transform=val_transform)
    return train_dst, val_dst


def validate(opts, model, loader, device, metrics, epoch, ret_samples_ids=None):
    model.eval()
    tmp_path = os.path.join('tmp_results', TIME_NOW, str(epoch))
    os.makedirs(tmp_path)
    save_path = os.path.join('results', TIME_NOW, str(epoch))
    os.makedirs(save_path)
    denorm = utils.Denormalize(mean=[0.485, 0.456, 0.406],
                               std=[0.229, 0.224, 0.225])

    with torch.no_grad():
        for img_ids, images, labels, Ws, Hs in loader:
            images = images.to(device, dtype=torch.float32)
            labels = labels.to(device, dtype=torch.long)

            outputs = model(images)['out']
            preds = outputs.detach().max(dim=1)[1].cpu().numpy()
            targets = labels.cpu().numpy()
            for img_id, pred, image, target, w, h in zip(img_ids, preds, images, targets, Ws, Hs):
                image = image.detach().cpu().numpy()
                # target = targets[i]
                # pred = preds[i]
                image = (denorm(image) * 255).transpose(1, 2, 0).astype(np.uint8)
                target = loader.dataset.decode_target(target).astype(np.uint8)
                # pred = loader.dataset.decode_target(pred).astype(np.uint8)
                pred1 = loader.dataset.decode_target(pred).astype(np.uint8)
                pred2 = pred.astype(np.uint8)

                output_transform = transforms.Compose([
                    transforms.Resize([h, w]),
                ])
                output_transform(Image.fromarray(pred1)).save(tmp_path + '/{id}.png'.format(id=img_id))
                output_transform(Image.fromarray(target)).save(tmp_path + '/{id}_target.png'.format(id=img_id))
                output_transform(Image.fromarray(pred2).convert('P')).save(save_path + '/{id}.png'.format(id=img_id))

                fig = plt.figure()
                plt.imshow(output_transform(Image.fromarray(image)))
                plt.axis('off')
                plt.imshow(output_transform(Image.fromarray(pred.astype(np.uint8))), alpha=0.7)
                ax = plt.gca()
                ax.xaxis.set_major_locator(matplotlib.ticker.NullLocator())
                ax.yaxis.set_major_locator(matplotlib.ticker.NullLocator())
                plt.savefig(tmp_path + '/{id}_overlay.png'.format(id=img_id), bbox_inches='tight', pad_inches=0)
                plt.close()
        # eval_loader = Dataset_Seg(save_path, 'datasets/data/VOCdevkit/VOC2012/SegmentationClass', split='val').get_loader()
        eval_loader = Dataset_Seg(save_path, opts.anno_root,
                                  split='val').get_loader()
        miou, pacc, macc = evaluate(eval_loader, 21)
    return miou, pacc, macc


def main():
    opts = get_argparser().parse_args()
    if opts.dataset.lower() in ['voc', 'cv']:
        opts.num_classes = 21
    elif opts.dataset.lower() == 'cityscapes':
        opts.num_classes = 19
    # Setup visualization
    vis = Visualizer(port=opts.vis_port,
                     env=opts.vis_env) if opts.enable_vis else None
    if vis is not None:  # display options
        vis.vis_table("Options", vars(opts))

    os.environ['CUDA_VISIBLE_DEVICES'] = opts.gpu_id
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print("Device: %s" % device)

    # Setup random seed
    torch.manual_seed(opts.random_seed)
    np.random.seed(opts.random_seed)
    random.seed(opts.random_seed)

    train_dst, val_dst = get_dataset(opts)
    train_loader = data.DataLoader(
        train_dst, batch_size=opts.batch_size, shuffle=True, num_workers=1)
    val_loader = data.DataLoader(
        val_dst, batch_size=opts.val_batch_size, shuffle=True, num_workers=1)
    # train_loader = train_dst
    # val_loader = val_dst
    print("Train set: %d, Val set: %d" % (len(train_dst), len(val_dst)))

    # Set up model
    # model_map = {
    #     'deeplabv3_resnet50': network.deeplabv3_resnet50,
    #     'deeplabv3plus_resnet50': network.deeplabv3plus_resnet50,
    #     'deeplabv3_resnet101': network.deeplabv3_resnet101,
    #     'deeplabv3plus_resnet101': network.deeplabv3plus_resnet101,
    #     'deeplabv3_mobilenet': network.deeplabv3_mobilenet,
    #     'deeplabv3plus_mobilenet': network.deeplabv3plus_mobilenet
    # }
    #
    # model = model_map[opts.model](num_classes=opts.num_classes, output_stride=opts.output_stride)
    model = deeplabv3_resnet101(True, num_classes=21)

    def save_ckpt(path):
        """ save current model
        """
        torch.save(model.state_dict(), path)
        print("Model saved as %s" % path)

    utils.mkdir('checkpoints')
    # Restore
    # best_score = 0.0
    # cur_itrs = 0
    # cur_epochs = 0
    if opts.ckpt is not None and os.path.isfile(opts.ckpt):
        # https://github.com/VainF/DeepLabV3Plus-Pytorch/issues/8#issuecomment-605601402, @PytaichukBohdan
        # checkpoint = torch.load(opts.ckpt, map_location=torch.device('cpu'))
        checkpoint = torch.load(opts.ckpt)
        model.load_state_dict(checkpoint)
        # model.load_state_dict(checkpoint)
        # model = nn.DataParallel(model)
        model.to(device)
        # if opts.continue_training:
        #     optimizer.load_state_dict(checkpoint["optimizer_state"])
        #     scheduler.load_state_dict(checkpoint["scheduler_state"])
        #     cur_itrs = checkpoint["cur_itrs"]
        #     best_score = checkpoint['best_score']
        #     print("Training state restored from %s" % opts.ckpt)
        print("Model restored from %s" % opts.ckpt)
        del checkpoint  # free memory
    else:
        print("[!] Pytorch Pretrain")
        # model = nn.DataParallel(model)
        model.to(device)

    # if opts.separable_conv and 'plus' in opts.model:
    #     network.convert_to_separable_conv(model.classifier)
    # model = deeplabv3_resnet101(True, num_classes=21)
    # params = torch.load('')
    utils.set_bn_momentum(model.backbone, momentum=0.01)
    # Set up metrics
    metrics = StreamSegMetrics(opts.num_classes)

    # Set up optimizer
    optimizer = torch.optim.SGD(params=[
        {'params': model.backbone.parameters(), 'lr': 0.1 * opts.lr},
        {'params': model.classifier.parameters(), 'lr': opts.lr},
    ], lr=opts.lr, momentum=0.9, weight_decay=opts.weight_decay)
    if opts.lr_policy == 'poly':
        scheduler = utils.PolyLR(optimizer, opts.total_itrs, power=0.9)
    elif opts.lr_policy == 'step':
        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=opts.step_size, gamma=0.1)

    # Set up criterion
    if opts.loss_type == 'focal_loss':
        criterion = utils.FocalLoss(ignore_index=255, size_average=True)
    elif opts.loss_type == 'cross_entropy':
        criterion = nn.CrossEntropyLoss(ignore_index=255, reduction='mean')

    # ==========   Train Loop   ==========#
    vis_sample_id = np.random.randint(0, len(val_loader), opts.vis_num_samples,
                                      np.int32) if opts.enable_vis else None  # sample idxs for visualization
    denorm = utils.Denormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # denormalization for ori images

    if opts.test_only:
        model.eval()
        miou, pacc, macc = validate(opts=opts, model=model, loader=val_loader, device=device, metrics=metrics,
                                    ret_samples_ids=vis_sample_id)
        # print(metrics.to_str(val_score))
        return

    best_pacc, best_macc, best_miou = 0, 0, 0
    best_pacc_iter, best_macc_iter, best_miou_iter = 0, 0, 0

    checkpoint_path = os.path.join('checkpoints', opts.dataset, TIME_NOW)
    # create checkpoint folder to save model
    if not os.path.exists(checkpoint_path):
        os.makedirs(checkpoint_path)
    pacc, macc, miou = validate(opts=opts, model=model, loader=val_loader, device=device, metrics=metrics,
                                epoch=0, ret_samples_ids=vis_sample_id)
    if vis is not None:
        vis.vis_scalar("[Val] Pixel Acc", 0, pacc)
        vis.vis_scalar("[Val] Mean IoU", 0, miou)
        vis.vis_scalar("[Val] Mean Acc", 0, macc)
    for cur_iter in range(1, opts.total_itrs + 1):  # cur_itrs < opts.total_itrs:
        # =====  Train  =====
        model.train()
        epoch_loss = 0
        for batch_id, (img_ids, images, labels, w, h) in enumerate(train_loader):
            # print('ids:' + ' '.join(img_ids))
            # images = Variable(torch.unsqueeze(images, dim=0).float(), requires_grad=False)
            # labels = Variable(torch.unsqueeze(labels, dim=0).float(), requires_grad=False)
            images = images.to(device, dtype=torch.float32)
            labels = labels.to(device, dtype=torch.long)
            optimizer.zero_grad()
            outputs = model(images)['out']
            # outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            np_loss = loss.detach().cpu().numpy()
            epoch_loss += np_loss
            print('Training Epoch: {epoch} [{trained_samples}/{total_samples}]\tLoss: {:0.4f}\tLR: {:0.6f}'.format(
                np_loss,
                optimizer.param_groups[0]['lr'],
                epoch=cur_iter,
                trained_samples=batch_id * opts.batch_size + len(images),
                total_samples=len(train_loader.dataset)
                # total_samples=len(train_loader)
            ))
        if vis is not None:
            vis.vis_scalar('Loss', cur_iter, epoch_loss)
        # print("Epoch {epoch}, Loss={loss}".format(epoch=cur_iter, loss=epoch_loss / len(train_loader.dataset)))
        # print("Epoch {epoch}, Loss={loss}".format(epoch=cur_iter, loss=epoch_loss / len(train_loader)))
        print("validation...")
        pacc, macc, miou = validate(opts=opts, model=model, loader=val_loader, device=device, metrics=metrics,
                                    epoch=cur_iter, ret_samples_ids=vis_sample_id)
        to_save = False
        if pacc > best_pacc:
            best_pacc = pacc
            best_pacc_iter = cur_iter
            to_save = True
        if macc > best_macc:
            best_macc = macc
            best_macc_iter = cur_iter
            to_save = True
        if miou > best_miou:
            best_miou = miou
            best_miou_iter = cur_iter
            to_save = True
        if to_save:
            cur_checkpoint_path = os.path.join(checkpoint_path, '{epoch}-{:.4f}-{:.4f}-{:.4f}.pth')
            save_ckpt(os.path.join(cur_checkpoint_path.format(macc, miou, pacc, epoch=cur_iter)))
        if vis is not None:
            vis.vis_scalar("[Val] Pixel Acc", cur_iter, pacc)
            vis.vis_scalar("[Val] Mean IoU", cur_iter, miou)
            vis.vis_scalar("[Val] Mean Acc", cur_iter, macc)
        print("best macc, miou, pacc: {macc}({i_macc}), {miou}({i_miou}), {pacc}({i_pacc})"
              .format(pacc=best_pacc, i_pacc=best_pacc_iter,
                      macc=best_macc, i_macc=best_macc_iter,
                      miou=best_miou, i_miou=best_miou_iter))
        scheduler.step()


if __name__ == '__main__':
    main()
